__author__ = 'Bill Liu'

import copy
import math
from matrix import *


def get_non_basis(var, basis):
    non_basis = []
    for i in var:
        if i not in basis:
            non_basis.append(i)
    return non_basis

def disp_step(mat, basis, selected=None):
    println = mat.format_print(bracket=False, hold_on=True, selected=selected)
    n_var = mat.get_row_num() - 1
    var_width = int(math.ceil(math.log10(n_var)))
    mat_width = len(println[0])
    n_print = len(println)
    split_pos = var_width + 2 + mat_width - 4 - mat.get_col_width(-1)

    for i in range(n_print):
        if i == n_print - 3:
            println[i] = ' '*(var_width + 2) + '-' * mat_width
        else:
            if i % 2 == 1 and i != n_print - 2:
                println[i] = 'x' + str(basis[i/2]) + ' ' + println[i]
            else:
                println[i] = ' '*(var_width + 2) + println[i]

            if i != 0 and i != n_print - 1:
                println[i] = println[i][:split_pos] + '|' + println[i][split_pos+1:]

    for line in println:
        print(line)

def check_unique_solution(mat, basis):
    n_var = mat.get_col_num() - 1
    var = range(n_var)
    non_basis = get_non_basis(var, basis)
    mat_test = mat.get_row(-1).get_col_slice(non_basis).get_data()[0]
    unique_solution = True
    idx = -1
    for i in range(len(mat_test)):
        if mat_test[i] == 0:
            unique_solution = False
            idx = i
            break
    return unique_solution, idx

def find_exchange(mat, var, basis):
    non_basis = get_non_basis(var, basis)
    if mat.get_row(-1).get_col_slice(non_basis).element_negative():
        return basis, non_basis, False, True
    else:
        mat_test = mat.get_row(-1).get_col_slice(non_basis).get_data()[0]
        max_val = -1e6
        for i in range(len(mat_test)):
            if mat_test[i] > max_val:
                max_val = mat_test[i]
                idx = i
        in_basis = non_basis[idx]

        mat_test = mat.get_col(in_basis).get_row_slice(range(len(basis)))
        no_solution = mat_test.element_negative()

        b = mat.get_col(-1).get_row_slice(range(len(basis)))
        mat_test = divide(b, mat_test).get_data()
        min_val = 1e6
        for i in range(len(mat_test)):
            if 0 <= mat_test[i][0] < min_val:
                min_val = mat_test[i][0]
                idx = i
        out_basis = basis[idx]

        return in_basis, out_basis, no_solution, False

def get_result(tabular, var, basis):
        target_val = tabular.get_element(-1, -1)
        b = tabular.get_col(-1).get_row_slice(range(len(basis))).get_data()

        var_val = []
        for i in var:
            if i not in basis:
                var_val.append(0)
            else:
                var_val.append(b[basis.index(i)][0])
        return target_val, var_val

def optimize_by_identity(constraint_matrix, B_inv_b, target_func, stop_criteria='opt', out_of_basis=None, basis=None, basis_pos=None):
    if stop_criteria != 'opt' and stop_criteria != 'out_of_basis':
        print('Stop criteria should be \'opt\' or \'out_of_basis\'')

    if basis is None or basis_pos is None:
        basis, basis_pos = constraint_matrix.find_identity()
        n_basis = len(basis)
        if n_basis < constraint_matrix.get_row_num():
            print('The function \'optimize_by_identity\' only support matrix contains identity matrix')
            return None

    n_var = constraint_matrix.get_col_num()
    var = range(n_var)

    assert(n_var == target_func.get_col_num())
    main_matrix = copy.deepcopy(constraint_matrix)           # the whole tabular is constructed in a single matrix.
    main_matrix.hstack(B_inv_b)                             # to prohibit memory corruption, copy the original matrix first
    main_matrix.reformat(basis_pos)                         # for beautiful display, arrange the identity matrix to normal order
    # rows have been re-ordered, so the constraints should be changed!
    constraint_matrix = main_matrix.get_row_slice(range(constraint_matrix.get_row_num())).get_col_slice(range(n_var))
    B_inv_b = main_matrix.get_row_slice(range(constraint_matrix.get_row_num())).get_col(-1)

    non_basis = get_non_basis(var, basis)
    B_inv_P = constraint_matrix.get_col_slice(non_basis)
    c_B = target_func.get_col_slice(basis)
    c_N = target_func.get_col_slice(non_basis)
    test_num = dot(c_B, B_inv_P)
    test_num.subtract(c_N)
    test_num = test_num.get_data()                         # get testing number, using c_b * (B^-1) * N - c_N

    lower_row = zeros(1, n_var)
    for i in range(n_var):
        if i in basis:
            lower_row.set_element(row=0, col=i, val=0)
        else:
            lower_row.set_element(row=0, col=i, val=test_num[0][non_basis.index(i)])
    func = dot(c_B, B_inv_b)
    lower_row.hstack(func)
    main_matrix.vstack(lower_row)                            # construct the complete tabular

    while True:
        in_basis, out_basis, no_solution, found_solution = \
            find_exchange(main_matrix, var, basis)

        # test stop criteria
        if stop_criteria == 'out_of_basis':
            is_out_of_basis = True
            for i in out_of_basis:
                if i in basis:
                    is_out_of_basis = False
                    break
            if is_out_of_basis is True:
                break
        if no_solution or found_solution:                    # check normal criteria
            break

        main_row = basis.index(out_basis)
        disp_step(main_matrix, basis, selected=[main_row, in_basis])

        # replace the basis with in_basis
        basis[main_row] = in_basis
        main_matrix.gaussian_elimination(main_row, in_basis)

    disp_step(main_matrix, basis)                             # show the final matrix

    # whether we have different task, the solution MUST exist
    # when dealing with the result, we should check no_solution to be False first
    if no_solution:
        return main_matrix, no_solution, None, None, None, None

    target_val, var_val = get_result(main_matrix, var, basis)

    if stop_criteria == 'out_of_basis':
        if not is_out_of_basis:
            # the artificial variables are still basis, so check their values
            is_art_basis_zeros = True
            art_basis_zeros = []
            for i in out_of_basis:
                if i in basis:
                    art_basis_zeros.append(i)
                    if var_val[i] != 0:
                        is_art_basis_zeros = False
                        break
            if is_art_basis_zeros:
                origin_var = []
                for i in var:
                    if i not in out_of_basis:
                        origin_var.append(i)
                # x_a = 0, get the artificial variable out of basis
                # 1. replace the artificial variables with original
                # 2. or the original coefficients in this row are all zeros, delete
                for art in art_basis_zeros:
                    out_row = basis.index(art)
                    out_row_data = main_matrix.get_row(out_row).get_data()[0]
                    # check original coefficients
                    origin_all_zeros = True
                    for i in origin_var:
                        if out_row_data[i] != 0:
                            origin_all_zeros = False
                            break
                    if origin_all_zeros:
                        # delete out_row
                        new_main_matrix = main_matrix.get_row_slice(range(out_row))
                        new_main_matrix.vstack(main_matrix.get_row_slice(range(out_row+1, main_matrix.get_row_num())))
                        main_matrix = new_main_matrix
                        del basis[out_row]
                    else:
                        # replace artificial coefficient
                        in_basis = -1
                        for i in var:
                            if i not in out_of_basis and i not in basis:
                                in_basis = i
                                break
                        disp_step(main_matrix, basis, selected=[out_row, in_basis])                # swap the basis
                        basis[out_row] = in_basis
                        main_matrix.gaussian_elimination(out_row, in_basis)
                is_out_of_basis = True
                disp_step(main_matrix, basis)
        return main_matrix, no_solution, True, target_val, target_val, is_out_of_basis
    else:
        is_unique, zero_non_basis = check_unique_solution(main_matrix, basis)
        return main_matrix, no_solution, is_unique, target_val, var_val, None

class Simplex:
    def __init__(self, target_function=None, constraint_A=None, constraint_b=None, disp_style=None):
        self.target_function = target_function
        self.constraint_A = constraint_A
        self.constraint_b = constraint_b
        self.disp_style = disp_style

        if self.check_prob() is False:
            self.reset()

    def update_prob(self, target_function=None, constraint_A=None, constraint_b=None, disp_style=None):
        self.target_function = target_function
        self.constraint_A = constraint_A
        self.constraint_b = constraint_b
        self.disp_style = disp_style

        if self.check_prob() is False:
            self.reset()

    def reset(self):
        self.target_function = None
        self.constraint_A = None
        self.constraint_b = None
        self.disp_style = None

    def check_prob(self):
        check_ok = True
        if self.target_function.get_row_num() != 1:
            print('Target function must be a row vector')
            check_ok = False
        if self.constraint_b.get_col_num() != 1:
            print('constrain b must be a column vector')
            check_ok = False
        if self.target_function.get_col_num() != self.constraint_A.get_col_num():
            print('The variable in target function and constraint A do not match')
            check_ok = False
        if self.constraint_A.get_row_num() != self.constraint_b.get_row_num():
            print('The number of constraint in A and b do not match')
            check_ok = False
        if self.constraint_b.element_positive() is False:
            print('The constraint b has negative elements')
            check_ok = False
        return check_ok

    def disp_prob(self):
        print('The problem is:\nTarget function:\nmin'),
        idx = 0
        n_var = len(self.target_function.get_data()[0])

        func = ''
        for x in self.target_function.get_data()[0]:
            if x == 0:
                continue
            if x > 0:
                if idx > 0:
                    func += '+'
                else:
                    func += ' '
            if x == 1 or x == -1:
                if x == -1:
                    func += '-'
            else:
                func += '{0}'.format(x)
            func += 'x{0}'.format(idx)
            idx += 1
        print(func),
        print('\ns.t.')

        # constrain matrix
        println_A = self.constraint_A.format_print(bracket=True, hold_on=True)

        # variables
        println_x = []
        width = int(math.ceil(math.log10(n_var))) + 1
        line = ' _'
        for i in range(width):
            line += ' '
        line += '_ '
        println_x.append(line)
        for i in range(n_var):
            line = '| x{0} |'.format(i)
            println_x.append(line)
            if i != n_var - 1:
                line = '| '
                for j in range(width):
                    line += ' '
                line += ' |'
            else:
                line = ' -'
                for i in range(width):
                    line += ' '
                line += '- '
            println_x.append(line)

        # constrain b
        println_b = self.constraint_b.format_print(bracket=True, hold_on=True)
        assert(len(println_A) == len(println_b))

        # combine all the matrix
        println = []
        n_row_constraint = len(println_A)
        n_row_var = len(println_x)
        n_width_A = len(println_A[0])
        n_width_x = len(println_x[0])
        for i in range(n_row_constraint):
            println.append(println_A[i] + println_x[i] + '   ' + println_b[i])
        for i in range(n_row_var - n_row_constraint):
            line = ''
            for j in range(n_width_A):
                line += ' '
            line += println_x[n_row_constraint + i]
            println.append(line)

        # add '+'
        row_plus = n_row_constraint / 2
        col_plus = n_width_A + n_width_x + 1
        println[row_plus] = println[row_plus][:col_plus] + '=' + println[row_plus][col_plus+1:]

        for line in println:
            print(line)
        print('Note: all variable should be positive!!! x > 0')

    def solve_prob(self):
        if self.target_function is None or self.constraint_A is None or self.constraint_b is None \
                or self.disp_style is None:
            print('Incomplete problem description')
            return

        if self.disp_style != 'tabular':
            print('Display style is not supported')
            return

        self.disp_prob()

        basis, basis_pos = self.constraint_A.find_identity()
        n_basis = len(basis)
        n_var = self.constraint_A.get_col_num()
        n_rank = self.constraint_A.get_row_num()
        var = range(n_var)                                                   # TODO: the variables' name could be different in optimize_by_identity! try to fix it?
        if n_basis < n_rank:
            print('Use Two Phase Method\nImport artificial variables')

            # Phase I: add artificial variables and optimize the problem
            n_art = n_rank - n_basis                # artificial variables needed to construct identity matrix
            n_var_all = n_var + n_art
            var_art = range(n_var, n_var_all)
            var_all = var + var_art
            print('Import {0} variables: '.format(n_art)),
            for i in var_art:
                print('x{0}, '.format(i)),
            print('to construct identity matrix in simplex tabular')

            # create new target function
            art_target = zeros(1, n_var)
            target_ones = ones(1, n_art)
            art_target.hstack(target_ones)

            # create new constraint
            art_constraint_b = copy.deepcopy(self.constraint_b)

            art_constraint_A_add = zeros(n_rank, n_art)
            idx_art = 0
            for i in range(n_rank):
                if i not in basis_pos:
                    basis_pos.append(i)
                    basis.append(var_art[idx_art])
                    art_constraint_A_add.set_element(i, var_art[idx_art] - n_var, 1)
                    idx_art += 1
            art_constraint_A = copy.deepcopy(self.constraint_A)
            art_constraint_A.hstack(art_constraint_A_add)
            main_matrix, no_solution, is_unique, target_val, var_val, is_out_of_basis =\
                optimize_by_identity(art_constraint_A, art_constraint_b, art_target, stop_criteria='out_of_basis', out_of_basis=var_art, basis=basis, basis_pos=basis_pos)

            # n_rank should be updated, because some rows may be deleted in Phase i
            n_rank = main_matrix.get_row_num() - 1

            if no_solution or not is_out_of_basis:
                print('The problem has no solutions')
                return

            # Phase II: eliminate the artificial variables and solve the original problem
            self.constraint_A = main_matrix.get_row_slice(range(n_rank)).get_col_slice(var)
            self.constraint_b = main_matrix.get_row_slice(range(n_rank)).get_col(-1)
            basis, basis_pos = self.constraint_A.find_identity()

        # Gaussian Elimination
        print('Choose: '),
        for x in basis:
            print('x%d, ' % x),
        print('as initial basis. We use Gaussian Elimination Algorithm')

        dummy, no_solution, is_unique, target_val, var_val, dummy =\
            optimize_by_identity(self.constraint_A, self.constraint_b, self.target_function, stop_criteria='opt', basis=basis, basis_pos=basis_pos)

        if no_solution is True:
            print('The problem is not bounded!')
        else:
            print('The problem has solved.\nThe minimum target value is {0}\nAchieved at ('.format(target_val)),
            for i in var:
                print('x{0}'.format(i)),
            print(') = ('),
            for i in var_val:
                print('{0},'.format(i)),
            print(')')
            if not is_unique:
                print('The problem has multiple solutions')







